import math
import pickle
import numpy
from scipy import interpolate
from scipy.optimize import minimize
from scipy.optimize import curve_fit


def PolyTwoGauss(p, x):           
        # offset = distance between peaks
        # p[0], p[1], p[2] = amplitude , center_x, sigma: for peak A
        # p[3], p[4] = amplitude, sigma: for peak B
        # p[5] = polynomial offset
        # p[6], p[7], p[8] = ax^2 + bx + c
    offset = 52
    y = numpy.zeros(len(x[0]))
    for i, j in enumerate(x[0]):
        y[i] = p[0]*math.exp(-0.5*((j-p[1])/p[2])**2.) + p[3]*math.exp(-0.5*((j-(p[1]+offset))/p[4])**2.) + p[6]*(j - p[5])**2 + p[7]*(j - p[5]) + p[8]
        if y[i] <= 1e-5: 
            y[i] = 1e-5
    return y


def PolyFn(x, x0, a, b, c):
   return a*(x-x0)**2 + b*(x-x0) + c


def Residual(p, x, y):
    modely = PolyTwoGauss(p, x)
    return numpy.sum((modely - y)**2)


def BinData(x, y):
    bins = 200.0
    max_d = 400.0
    xdata = numpy.linspace(0, max_d, bins)
    ydata = numpy.zeros(bins)

    for j, k in enumerate(xdata):
        pos_torque = []
        x0 = k - (max_d / bins) / 2.
        x1 = k + (max_d / bins) / 2.
        for i, p in enumerate(x):
           if p > x0 and p <= x1:
              pos_torque.append(y[i])
        if len(pos_torque) > 0:
            ydata[j] += max(pos_torque)
    
    return xdata, ydata


def CorrectForOffset_G(x, y, station):
    new_pos = []
    if x[0] != 0: 
        p0 = x[0]
        for j in x:
            new_pos.append(j - p0)
    else: 
        new_pos = x
    
    xdata, ydata = BinData(new_pos, y)
    
    if station == 'S2':
        xbin = xdata[140:170]
        ybin = ydata[140:170]
    else:
        xbin = xdata[140:-10]
        ybin = ydata[140:-10]
    
    poly_fit = curve_fit(PolyFn, xbin, ybin) #to get initial guess on the baseline
    poly_parms = poly_fit[0]
    
    fit_parms = [5, 315, 5, 5, 5, poly_parms[0], poly_parms[1], poly_parms[2], poly_parms[3]]
    fit_bounds = [(2, 20), (280, 330), (1, 10), (2, 20), (1, 10), (None, None), (None,None), (None, None), (None, None)]
    result = minimize(Residual, fit_parms, args = (([xbin]), ([ybin])), method='TNC', bounds=fit_bounds,tol=0.0000005,options={'disp':False})
    new_parms = result.x
    
    offset = (new_parms[1] - fit_parms[1]) + x[0]
    
    # produces position data array without the offset
    offsetfix = []
    for i, k in enumerate(x):
        newx = k - offset
        offsetfix.append(newx)
    
    return offsetfix, offset, new_parms
